//
// Updated by 王岩松 on 2020/3/23.
//
// 修改自(https://github.com/lesterphillip/polynomial-root-searcher/blob/5d13ce8f391fc02f9f57323ba5610af2a4393bc8/root-solver/main.cpp)

#ifndef FORMANT_ROOTS_H
#define FORMANT_ROOTS_H

#include <complex>
#include <cmath>
#include <vector>

#define  ROOTS_ITERATER_LIMIT 10000

template <class T>
static std::complex<T> cal_poly(T poly[], int deg, std::complex<T> x) {
    std::complex<T> y = 0;
    for(int i = 0 ; i <= deg; ++ i)
        y += poly[i] * pow(x,(deg-i));
    return y;
}

template <class T>
static std::complex<T> cal_poly(std::vector<T> &poly, int deg, std::complex<T> x) {
    std::complex<T> y = 0;
    for(int i = 0 ; i <= deg; ++ i)
        y += poly[i] * pow(x,(deg-i));
    return y;
}

// poly[0] x^deg + poly[1] x^(deg-1) + ... + poly[deg], require poly[deg] = 1
template <class T>
std::complex<T> *roots(T poly[], int deg) {
    std::complex<T> constant(0.4,0.9), reset(1,0);
    std::complex<T> *rootApprox = NULL;
    std::complex<T> firstDenom[deg], secondDenom[deg];
    rootApprox = new std::complex <T> [deg];
    fill_n(firstDenom, deg, 1); //replaces all the values of the array with 1
    fill_n(secondDenom, deg, 1); //replaces all the values of the array with 1
    for (int i=0; i<deg; i++)
        rootApprox[i] = pow(constant, i);
    //iterates for a certain number of times until the root...
    //converges to the real number
    for (int iterations=0; iterations<=ROOTS_ITERATER_LIMIT; iterations++) {
        //for loop goes through all the roots available
        //highest degree states the number of roots needed
        for(int nthRoot = 0; nthRoot<deg; nthRoot++) {
            //for the initializer
            //computes for the part of pn-1 - qn-1 and so on..
            for(int firstHalf=1; firstHalf<=deg-1; firstHalf++) {
                //ensures that the array starts at 1+0i for the first computation
                if(firstHalf==1)
                    firstDenom[nthRoot]=reset;
                if(nthRoot<firstHalf)
                    firstDenom[nthRoot] *= (rootApprox[nthRoot] - rootApprox [firstHalf]);
            }
            //initializes the polynomial function
            //f(qn-1), f(pn-1) ... etc
            std::complex <T> y = cal_poly(poly, deg, rootApprox[nthRoot]);
            //computes for  (qn-1 - pn) , (rn-1 - pn), ... etc
            //starts at nthRoot > 0 because formula says so
            if(nthRoot>0) {
                for(int secondHalf=0; secondHalf<=deg-1; secondHalf++) {
                    //ensures that the array starts at 1+0i for the first computation
                    if(secondHalf==0)
                        secondDenom[nthRoot]=reset;
                    if(nthRoot>secondHalf)
                        secondDenom[nthRoot] *= (rootApprox[nthRoot] - rootApprox[secondHalf]);
                }
            }
            rootApprox[nthRoot] = (rootApprox [nthRoot] - (y/(firstDenom[nthRoot]*secondDenom[nthRoot])));
        }
    }
    return rootApprox;
}

template <class T>
std::vector<std::complex<T>> roots(std::vector<T> &poly, int deg, std::vector<std::complex<T>> INITroot={}) {
    std::complex<T> constant(0.4,0.9), reset(1,0);
    std::vector<std::complex<T>> rootApprox;
    std::complex<T> firstDenom[deg], secondDenom[deg];
    rootApprox.resize(deg,0);

    fill_n(firstDenom, deg, 1); //replaces all the values of the array with 1
    fill_n(secondDenom, deg, 1); //replaces all the values of the array with 1
    for (int i=0; i<deg; i++)
        rootApprox[i] = pow(constant, i);
    if(INITroot.size() != 0) {
        for (int i=0; i<deg; i++)
            rootApprox[i] = INITroot[i];
    }
    //iterates for a certain number of times until the root...
    //converges to the real number
    for (int iterations=0; iterations<=ROOTS_ITERATER_LIMIT; iterations++) {
        long double SUMy = 0;
        //for loop goes through all the roots available
        //highest degree states the number of roots needed
        for(int nthRoot = 0; nthRoot<deg; nthRoot++) {
            //for the initializer
            //computes for the part of pn-1 - qn-1 and so on..
            for(int firstHalf=1; firstHalf<=deg-1; firstHalf++) {
                //ensures that the array starts at 1+0i for the first computation
                if(firstHalf==1)
                    firstDenom[nthRoot]=reset;
                if(nthRoot<firstHalf)
                    firstDenom[nthRoot] *= (rootApprox[nthRoot] - rootApprox [firstHalf]);
            }
            //initializes the polynomial function
            //f(qn-1), f(pn-1) ... etc
            std::complex <T> y = cal_poly(poly, deg, rootApprox[nthRoot]);
            SUMy += abs(y);
            //computes for  (qn-1 - pn) , (rn-1 - pn), ... etc
            //starts at nthRoot > 0 because formula says so
            if(nthRoot>0) {
                for(int secondHalf=0; secondHalf<=deg-1; secondHalf++) {
                    //ensures that the array starts at 1+0i for the first computation
                    if(secondHalf==0)
                        secondDenom[nthRoot]=reset;
                    if(nthRoot>secondHalf)
                        secondDenom[nthRoot] *= (rootApprox[nthRoot] - rootApprox[secondHalf]);
                }
            }
            rootApprox[nthRoot] = (rootApprox [nthRoot] - (y/(firstDenom[nthRoot]*secondDenom[nthRoot])));
        }
        if(SUMy < 1e-8) break;
    }
    return rootApprox;
}

template <class T>
void Roots_fixIntoUnitCircle (std::complex<T> *rt, int deg) {
    std::complex<T> z10 ( 1.0, 0.0 );
    for (int iroot = 0; iroot < deg; iroot ++)
        if (abs (rt [iroot]) > 1.0)
            rt [iroot] = z10 / std::conj (rt [iroot]);
}

template <class T>
void Roots_fixIntoUnitCircle (std::vector<std::complex<T>> & rt, int deg) {
    std::complex<T> z10 ( 1.0, 0.0 );
    for (int iroot = 0; iroot < deg; iroot ++)
        if (abs (rt [iroot]) > 1.0)
            rt [iroot] = z10 / std::conj (rt [iroot]);
}

#endif //FORMANT_ROOTS_H
